import paddle
from paddle import nn
from paddle.vision.models import resnet152
from paddle.vision import transforms 

import pdb

class EncoderCNN(paddle.nn.Layer):
    def __init__(self, embed_size=256): #batch_size):
        """Load the pretrained ResNet-152."""
        super(EncoderCNN, self).__init__()
        resnet = resnet152(pretrained=True)
        modules = list(resnet.children())[:-1]
        self.resnet = nn.Sequential(*modules)
        self.linear = nn.Linear(2048, embed_size)
        self.bn = nn.BatchNorm1D(embed_size, momentum=0.01)
    
    def forward(self, image_batch):
        """
        Extract feature vectors from input image batch
        """
        pdb.set_trace()
        features = self.resnet(image_batch)
        features = paddle.reshape(features, (-1, features[0].size))
        features = self.linear(features)
        features = self.bn(features)
        return features

